#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
#include <iostream>
#include <vector>
#include <algorithm>
#include <chrono>
#include "utiles.h"

// 自定义 Float16 结构体
struct Float16 {
    uint16_t value;

    // 从 float 转换为 Float16
    static Float16 FromFloat(float f) {
        Float16 result;
        uint32_t fbits = *reinterpret_cast<uint32_t*>(&f);
        uint32_t sign = (fbits >> 31) & 0x1;
        uint32_t exp = (fbits >> 23) & 0xff;
        uint32_t mant = fbits & 0x7fffff;

        // 处理特殊情况
        if (exp == 0xff) {
            // NaN 或 Inf
            result.value = (sign << 15) | 0x7c00;
            if (mant != 0) {
                // NaN
                result.value |= (mant & 0x3ff);
            }
        }
        else if (exp < 113) {
            // 下溢，置为 0
            result.value = sign << 15;
        }
        else if (exp > 142) {
            // 上溢，置为 Inf
            result.value = (sign << 15) | 0x7c00;
        }
        else {
            // 正常情况
            exp -= 112;
            mant >>= 13;
            result.value = (sign << 15) | (exp << 10) | mant;
        }

        return result;
    }

    // 从 Float16 转换为 float
    float ToFloat() const {
        uint32_t sign = (value >> 15) & 0x1;
        uint32_t exp = (value >> 10) & 0x1f;
        uint32_t mant = value & 0x3ff;

        if (exp == 0x1f) {
            // NaN 或 Inf
            uint32_t fbits = (sign << 31) | 0x7f800000;
            if (mant != 0) {
                // NaN
                fbits |= (mant << 13);
            }
            return *reinterpret_cast<float*>(&fbits);
        }
        else if (exp == 0) {
            // 非规范化数
            if (mant == 0) {
                return sign ? -0.0f : 0.0f;
            }
            while ((mant & 0x400) == 0) {
                mant <<= 1;
                exp--;
            }
            exp++;
            mant &= 0x3ff;
        }

        exp += 112;
        mant <<= 13;
        uint32_t fbits = (sign << 31) | (exp << 23) | mant;
        return *reinterpret_cast<float*>(&fbits);
    }
};

// 特化 Ort::TypeToTensorType 模板类
namespace Ort {
    template<>
    struct TypeToTensorType<Float16> {
        static constexpr ONNXTensorElementDataType type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
    };
}

// 定义类别
const std::vector<std::string> classes = { "Cavity", "discoloration", "Normal", "gingivitis" };

// 定义颜色调色板
std::vector<cv::Scalar> color_palette;

// 计算交并比（IoU）
float calculate_iou(const cv::Rect& box1, const cv::Rect& box2) {
    int x1 = std::max(box1.x, box2.x);
    int y1 = std::max(box1.y, box2.y);
    int x2 = std::min(box1.x + box1.width, box2.x + box2.width);
    int y2 = std::min(box1.y + box1.height, box2.y + box2.height);

    int intersection_area = std::max(0, x2 - x1) * std::max(0, y2 - y1);
    int box1_area = box1.width * box1.height;
    int box2_area = box2.width * box2.height;

    float iou = static_cast<float>(intersection_area) / (box1_area + box2_area - intersection_area);
    return iou;
}

// 自定义非最大抑制（NMS）
std::vector<int> custom_NMSBoxes(const std::vector<cv::Rect>& boxes, const std::vector<float>& scores, float confidence_threshold, float iou_threshold) {
    std::vector<int> indices;
    if (boxes.empty()) return indices;

    std::vector<int> sorted_indices(scores.size());
    for (size_t i = 0; i < scores.size(); ++i) {
        sorted_indices[i] = i;
    }
    // 按得分从高到低排序
    std::sort(sorted_indices.begin(), sorted_indices.end(), [&scores](int i1, int i2) {
        return scores[i1] > scores[i2];
        });

    std::vector<bool> suppressed(scores.size(), false);
    for (size_t i = 0; i < sorted_indices.size(); ++i) {
        int index = sorted_indices[i];
        if (suppressed[index]) continue;
        if (scores[index] < confidence_threshold) continue;

        indices.push_back(index);
        for (size_t j = i + 1; j < sorted_indices.size(); ++j) {
            int other_index = sorted_indices[j];
            if (suppressed[other_index]) continue;
            if (calculate_iou(boxes[index], boxes[other_index]) > iou_threshold) {
                suppressed[other_index] = true;
            }
        }
    }

    return indices;
}

// 在图像上绘制检测结果
void draw_detections(cv::Mat& img, const cv::Rect& box, float score, int class_id) {
    cv::Scalar color = color_palette[class_id];
    cv::rectangle(img, box, color, 2);

    std::string label = classes[class_id] + ": " + std::to_string(score);
    int baseline = 0;
    cv::Size label_size = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);

    int label_x = box.x;
    int label_y = box.y - 10;
    if (label_y < label_size.height) {
        label_y = box.y + label_size.height + 10;
    }

    cv::rectangle(img, cv::Point(label_x, label_y - label_size.height), cv::Point(label_x + label_size.width, label_y + label_size.height), color, cv::FILLED);
    cv::putText(img, label, cv::Point(label_x, label_y), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0), 1, cv::LINE_AA);
}

void example() {
    // 初始化 ONNX Runtime
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntime");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    // 加载 ONNX 模型
    const ORTCHAR_T* model_path = L"half.onnx";

    // L"half.onnx"
    Ort::Session session(env, model_path, session_options);

    // 获取输入输出信息
    Ort::AllocatorWithDefaultOptions allocator;
    auto input_shape = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
    int input_channels = input_shape[1];
    int input_height = input_shape[2];
    int input_width = input_shape[3];
    std::cout << input_height << " " << input_width << std::endl;

    // 加载图像
    const char* image_path = "images/img_1.jpg";
    cv::Mat image = cv::imread(image_path);
    if (image.empty()) {
        std::cerr << "加载图像失败" << std::endl;
        return -1;
    }

    // 初始化颜色调色板
    for (size_t i = 0; i < classes.size(); ++i) {
        cv::Scalar color(rand() % 256, rand() % 256, rand() % 256);
        color_palette.push_back(color);
    }

    // 调整尺寸并归一化
    cv::Mat resized;
    cv::resize(image, resized, cv::Size(input_width, input_height));
    cv::cvtColor(resized, resized, cv::COLOR_BGR2RGB);
    resized.convertTo(resized, CV_32F, 1.0 / 255.0);

    // 转换为CHW格式并提取float数据
    std::vector<cv::Mat> channels;
    cv::split(resized, channels);
    std::vector<float> input_data_float;
    input_data_float.reserve(input_channels * input_height * input_width);
    for (const auto& channel : channels) {
        input_data_float.insert(input_data_float.end(), channel.begin<float>(), channel.end<float>());
    }

    // 将float数据转换为FP16格式
    std::vector<Float16> input_data_half;
    input_data_half.reserve(input_data_float.size());
    for (float f : input_data_float) {
        input_data_half.push_back(Float16::FromFloat(f));
    }

    // 创建FP16输入张量
    std::vector<int64_t> input_dim = { 1, input_channels, input_height, input_width };
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<Float16>(
        memory_info,
        input_data_half.data(),
        input_data_half.size(),
        input_dim.data(),
        input_dim.size()
    );

    // 准备输入输出名称
    Ort::AllocatedStringPtr input_name = session.GetInputNameAllocated(0, allocator);
    Ort::AllocatedStringPtr output_name = session.GetOutputNameAllocated(0, allocator);
    const char* input_names[] = { input_name.get() };
    const char* output_names[] = { output_name.get() };

    // 记录推理开始时间
    auto start_time = std::chrono::high_resolution_clock::now();

    // 执行推理
    std::vector<Ort::Value> outputs = session.Run(
        Ort::RunOptions{ nullptr },
        input_names,
        &input_tensor,
        1,
        output_names,
        1
    );

    // 记录推理结束时间
    auto end_time = std::chrono::high_resolution_clock::now();

    // 计算推理时间（毫秒）
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
    std::cout << "推理时间: " << duration << " 毫秒" << std::endl;

    // 解析FP16输出数据
    Float16* output_data = outputs[0].GetTensorMutableData<Float16>();
    auto output_shape_info = outputs[0].GetTensorTypeAndShapeInfo();
    auto output_dims = output_shape_info.GetShape();
    const int num_classes = output_dims[1] - 4;
    const int num_boxes = output_dims[2];

    // 置信度阈值和IoU阈值
    const float confidence_thres = 0.1f;
    const float iou_thres = 0.5f;

    std::vector<cv::Rect> boxes;
    std::vector<float> scores;
    std::vector<int> class_ids;

    // 遍历所有检测框
    for (int i = 0; i < num_boxes; ++i) {
        // 提取位置信息并转换为float
        const float x = output_data[i].ToFloat();
        const float y = output_data[num_boxes + i].ToFloat();
        const float w = output_data[2 * num_boxes + i].ToFloat();
        const float h = output_data[3 * num_boxes + i].ToFloat();

        // 提取类别得分并找出最大值
        float max_score = -1.0f;
        int class_id = -1;
        for (int c = 0; c < num_classes; ++c) {
            const float score = output_data[(4 + c) * num_boxes + i].ToFloat();
            if (score > max_score) {
                max_score = score;
                class_id = c;
            }
        }

        // 输出符合条件的结果
        if (max_score >= confidence_thres) {
            // 计算边界框的缩放坐标
            float x_factor = static_cast<float>(image.cols) / input_width;
            float y_factor = static_cast<float>(image.rows) / input_height;
            int left = static_cast<int>((x - w / 2) * x_factor);
            int top = static_cast<int>((y - h / 2) * y_factor);
            int width = static_cast<int>(w * x_factor);
            int height = static_cast<int>(h * y_factor);

            boxes.emplace_back(left, top, width, height);
            scores.push_back(max_score);
            class_ids.push_back(class_id);
        }
    }

    // 应用非最大抑制过滤重叠的边界框
    std::vector<int> indices = custom_NMSBoxes(boxes, scores, confidence_thres, iou_thres);

    // 在图像上绘制非最大抑制后的边界框
    cv::Mat result_image = image.clone();
    for (int index : indices) {
        draw_detections(result_image, boxes[index], scores[index], class_ids[index]);
    }

    // 保存结果图像
    cv::imwrite("output_image.jpg", result_image);

    // 显示结果图像
    cv::imshow("Output", result_image);
    cv::waitKey(0);

    return 0;
}